package org.apache.flink.table.filesystem;//
// Source code recreated from a .class file by IntelliJ IDEA
// (powered by FernFlower decompiler)
//

import java.time.Duration;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.connector.file.table.PartitionFetcher;
import org.apache.flink.connector.file.table.PartitionReader;
import org.apache.flink.table.data.GenericRowData;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.functions.FunctionContext;
import org.apache.flink.table.functions.TableFunction;
import org.apache.flink.table.runtime.typeutils.InternalSerializers;
import org.apache.flink.table.runtime.typeutils.InternalTypeInfo;
import org.apache.flink.table.types.logical.RowType;
import org.apache.flink.util.FlinkRuntimeException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class FileSystemLookupFunction<P> extends TableFunction<RowData> {
    private static final Logger LOG = LoggerFactory.getLogger(FileSystemLookupFunction.class);
    private static final int MAX_RETRIES = 3;
    private static final Duration RETRY_INTERVAL = Duration.ofSeconds(10L);
    private final PartitionFetcher<P> partitionFetcher;
    private final PartitionFetcher.Context<P> fetcherContext;
    private final PartitionReader<P, RowData> partitionReader;
    private final RowData.FieldGetter[] lookupFieldGetters;
    private final Duration reloadInterval;
    private final TypeSerializer<RowData> serializer;
    private final RowType rowType;
    private transient Map<RowData, List<RowData>> cache;
    private transient long nextLoadTime;

    public FileSystemLookupFunction(PartitionFetcher<P> partitionFetcher, PartitionFetcher.Context<P> fetcherContext, PartitionReader<P, RowData> partitionReader, RowType rowType, int[] lookupKeys, Duration reloadInterval) {
        this.fetcherContext = fetcherContext;
        this.partitionFetcher = partitionFetcher;
        this.partitionReader = partitionReader;
        this.rowType = rowType;
        this.lookupFieldGetters = new RowData.FieldGetter[lookupKeys.length];

        for(int i = 0; i < lookupKeys.length; ++i) {
            this.lookupFieldGetters[i] = RowData.createFieldGetter(rowType.getTypeAt(lookupKeys[i]), lookupKeys[i]);
        }

        this.reloadInterval = reloadInterval;
        this.serializer = InternalSerializers.create(rowType);
    }

    public void open(FunctionContext context) throws Exception {
        super.open(context);
        this.cache = new HashMap();
        this.nextLoadTime = -1L;
        this.fetcherContext.open();
    }

    public TypeInformation<RowData> getResultType() {
        return InternalTypeInfo.of(this.rowType);
    }

    public void eval(Object... values) {
        this.checkCacheReload();
        RowData lookupKey = GenericRowData.of(values);
        List<RowData> matchedRows = (List)this.cache.get(lookupKey);
        if (matchedRows != null) {
            Iterator var4 = matchedRows.iterator();

            while(var4.hasNext()) {
                RowData matchedRow = (RowData)var4.next();
                this.collect(matchedRow);
            }
        }

    }

    private void checkCacheReload() {
        if (this.nextLoadTime <= System.currentTimeMillis()) {
            if (this.nextLoadTime > 0L) {
                LOG.info("Lookup join cache has expired after {} minute(s), reloading", this.reloadInterval.toMinutes());
            } else {
                LOG.info("Populating lookup join cache");
            }

            int numRetry = 0;

            while(true) {
                try {
                    long count = 0L;
                    GenericRowData reuse = new GenericRowData(this.rowType.getFieldCount());
                    this.partitionReader.open(this.partitionFetcher.fetch(this.fetcherContext));

                    RowData row;
                    while((row = (RowData)this.partitionReader.read(reuse)) != null) {
                        ++count;
                        RowData rowData = (RowData)this.serializer.copy(row);
                        RowData key = this.extractLookupKey(rowData);
                        List<RowData> rows = (List)this.cache.computeIfAbsent(key, (k) -> {
                            return new ArrayList();
                        });
                        rows.add(rowData);
                    }

                    this.partitionReader.close();
                    this.nextLoadTime = System.currentTimeMillis() + this.reloadInterval.toMillis();
                    LOG.info("Loaded {} row(s) into lookup join cache", count);
                    return;
                } catch (Exception var10) {
                    if (numRetry >= 3) {
                        throw new FlinkRuntimeException(String.format("Failed to load table into cache after %d retries", numRetry), var10);
                    }

                    ++numRetry;
                    long toSleep = (long)numRetry * RETRY_INTERVAL.toMillis();
                    LOG.warn(String.format("Failed to load table into cache, will retry in %d seconds", toSleep / 1000L), var10);

                    try {
                        Thread.sleep(toSleep);
                    } catch (InterruptedException var9) {
                        LOG.warn("Interrupted while waiting to retry failed cache load, aborting");
                        throw new FlinkRuntimeException(var9);
                    }
                }
            }
        }
    }

    private RowData extractLookupKey(RowData row) {
        GenericRowData key = new GenericRowData(this.lookupFieldGetters.length);

        for(int i = 0; i < this.lookupFieldGetters.length; ++i) {
            key.setField(i, this.lookupFieldGetters[i].getFieldOrNull(row));
        }

        return key;
    }

    public void close() throws Exception {
        this.fetcherContext.close();
    }

    @VisibleForTesting
    public Duration getReloadInterval() {
        return this.reloadInterval;
    }

    @VisibleForTesting
    public PartitionFetcher<P> getPartitionFetcher() {
        return this.partitionFetcher;
    }

    @VisibleForTesting
    public PartitionFetcher.Context<P> getFetcherContext() {
        return this.fetcherContext;
    }

    @VisibleForTesting
    public PartitionReader<P, RowData> getPartitionReader() {
        return this.partitionReader;
    }
}
